package com.jlf.starter.logger.reactive.decorator;

import com.google.gson.Gson;
import com.jlf.starter.logger.reactive.ApiLogger;
import com.jlf.starter.logger.reactive.entity.LogObject;
import com.jlf.starter.logger.reactive.provider.BodyProvider;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.lang.NonNull;
import org.springframework.util.FastByteArrayOutputStream;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.util.Map;

/**
 * 响应结果的装饰器
 *
 * @author wujr
 * 2023/8/3
 * @version 1.0
 * 【变更历史】
 * [wujr 2023/8/3 1.0 新增]
 */
@Slf4j
public class ResponseDecorator extends ServerHttpResponseDecorator {
    private final static Gson _gson = new Gson();
    private final BodyProvider bodyProvider = new BodyProvider();
    private final FastByteArrayOutputStream bodyOutputStream = new FastByteArrayOutputStream();

    @SuppressWarnings("unchecked")
    public ResponseDecorator(ServerWebExchange exchange, final LogObject logObject) {
        super(exchange.getResponse());
        getDelegate().beforeCommit(() -> {
            // 记录响应结果
            String bodyMessage;
            Map<String, Object> errors = (Map<String, Object>)exchange.getAttributes().get("__log_object");
            if (errors != null){
                exchange.getAttributes().remove("__log_object");
                bodyMessage = _gson.toJson(errors);
            }else {
                bodyMessage = bodyProvider.createBodyMessage(bodyOutputStream);
            }
            logObject.setResponseBody(bodyMessage);

            String error = getDelegate().getHeaders().getFirst("code");
            if (StringUtils.isBlank(error)){
                getDelegate().getHeaders().add("code", "0");
                logObject.setCode(0);
            }else{
                logObject.setCode(Integer.parseInt(error));
            }
            ApiLogger.log(logObject);

            return Mono.empty();
        });
    }
    @Override
    @NonNull
    public Mono<Void> writeWith(@NonNull Publisher<? extends DataBuffer> body) {
        Flux<DataBuffer> bodyBufferWrapper = Flux.from(body)
                .map(dataBuffer -> copyBodyBuffer(bodyOutputStream, dataBuffer));

        return super.writeWith(bodyBufferWrapper);
    }

    @Override
    @NonNull
    public Mono<Void> writeAndFlushWith(@NonNull Publisher<? extends Publisher<? extends DataBuffer>> body) {
        Flux<Flux<DataBuffer>> bodyBufferWrapper = Flux.from(body)
                .map(publisher -> Flux.from(publisher)
                        .map(buffer -> copyBodyBuffer(bodyOutputStream, buffer)));

        return super.writeAndFlushWith(bodyBufferWrapper);
    }


    private DataBuffer copyBodyBuffer(FastByteArrayOutputStream bodyStream, DataBuffer buffer) {
        try {
            ByteBuffer tmp = ByteBuffer.allocate(buffer.readableByteCount());
            buffer.toByteBuffer(tmp);
            Channels.newChannel(bodyStream).write(tmp.asReadOnlyBuffer());
            return buffer;
        } catch (IOException e) {
            log.error("复制body内容异常：", e);
            throw new RuntimeException(e);
        }
    }

    public static ServerHttpResponse build(ServerWebExchange exchange, LogObject logObject){
        // 对于下载文件的请求，不进行响应结果内容的记录
        MediaType contentType = exchange.getResponse().getHeaders().getContentType();
        if (contentType != null && contentType.includes(MediaType.MULTIPART_FORM_DATA)){
            // 这是文件下载，不记录响应内容
            exchange.getResponse()
                    .beforeCommit(() -> Mono.fromRunnable(() -> {
                        ApiLogger.log(logObject);
                    }));
            return exchange.getResponse();
        }else{
            // 需要把具体的内容记录到日志中
            return new ResponseDecorator(exchange, logObject);
        }
    }
}
